"""
yolo识别工具
"""

import cv2
import ultralytics
import numpy as np
import logging
logging.getLogger('ultralytics').setLevel(logging.WARNING)


class YoloDetector:
    model = None
    model_path = ""
    accepts_class_index = [
        0,
        1,
        2,
        3,
        5,
        6,
        7,
        9,
        10,
        11,
        12,
        13,
        14,
        15,
        16,
        17,
        18,
        19,
        24,
        25,
        26,
    ]

    def __init__(self, model_path):
        self.model_path = model_path
        self.model = ultralytics.YOLO(self.model_path, task="detect",verbose=False)

    def detect(self, img, depth_map=None):
        r = list()
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        hight, width, channel = img.shape
        results = self.model.predict(source=img, save=False, show=False)
        # print(results[0].names)
        # 解析预测结果
        for result in results:
            boxes = result.boxes
            for box in boxes:
                # 获取文本类别名称
                classname = result.names[int(box.cls.numpy()[0])]
                # 获取置信度
                confidence = box.conf.numpy()[0]
                # 获取边界框的坐标和尺寸
                cx, cy, w, h = box.xywhn.numpy()[0]
                # 创建字典并添加到结果列表
                if confidence < 0.5:
                    continue
                if int(box.cls.numpy()[0]) not in self.accepts_class_index:
                    continue
                if depth_map is not None:
                    x1, x2 = int(width * (cx - w / 2)), int(width * (cx + w / 2))
                    y1, y2 = int(hight * (cy - h / 2)), int(hight * (cy + h / 2))
                    distance = np.mean(depth_map[y1:y2, x1:x2])
                    r.append(
                        self.creatDict(classname, confidence, cx, cy, w, h, distance)
                    )
                else:
                    r.append(self.creatDict(classname, confidence, cx, cy, w, h))
        return r

    def creatDict(self, classname, confidence, cx, cy, w, h, distance="unknown"):
        return {
            "类名": classname,
            "置信度": float(confidence),
            "位置x": float(cx),
            "位置y": float(cy),
            "对象宽度w": float(w),
            "对象高度h": float(h),
            "百分比距离d": float(distance)/255,
        }


# 使用示例
if __name__ == "__main__":
    detector = YoloDetector(model_path="yolov8l.pt")
    image_path = "test.png"
    img = cv2.imread(image_path)
    detections = detector.detect(img)
    print(detections)
